from dotenv import load_dotenv, find_dotenv
from langchain_community.chat_message_histories import SQLChatMessageHistory
from langchain_core.messages import HumanMessage
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_openai import ChatOpenAI
from langchain.schema.output_parser import StrOutputParser

load_dotenv(find_dotenv())


def get_session_history(session_id):
    # 通过 session_id 区分对话历史，并存储在 sqlite 数据库中
    return SQLChatMessageHistory(session_id, "sqlite:///memory.db")


model = ChatOpenAI(model="deepseek-ai/DeepSeek-V3", temperature=0)

runnable = model | StrOutputParser()

runnable_with_history = RunnableWithMessageHistory(
    runnable, # 指定 runnable
    get_session_history, # 指定自定义的历史管理方法
)

runnable_with_history.invoke(
    [HumanMessage(content="你好，我叫王卓然")],
    config={"configurable": {"session_id": "wzr"}},
)

runnable_with_history.invoke(
    [HumanMessage(content="你知道我叫什么名字")],
    config={"configurable": {"session_id": "test"}},
)


# 获取指定 session_id 的历史记录
def print_history(session_id):
    history = get_session_history(session_id)
    print(f"\n=== Session ID: '{session_id}' 的对话历史 ===")
    for msg in history.messages:  # 遍历所有消息
        print(f"[{msg.type}] {msg.content}")  # 打印消息类型和内容

# 打印两个会话的历史
print_history("wzr")  # 你第一次调用的会话
print_history("test") # 第二次测试会话


